package com.koicarp.agent.image;

import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.util.ModelSerializer;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.ResNet50;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.File;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;

/**
 * @Author: liuxia
 * @CreateTime: 2025/3/4 下午7:53
 * @Description:
 */
public class ResNetFeatureExtractor {

    static {
        Path path = Paths.get(Core.NATIVE_LIBRARY_NAME);
        System.loadLibrary(path.toString()); }

    public static void main(String[] args) throws Exception {


        // 1. 自动下载并加载模型
        ComputationGraph model = ModelSerializer.restoreComputationGraph(
                new File("D:\\code\\langchain4j-agent-dev\\src\\main\\resources\\resnet50_dl4j_inference.v3.zip"));
//        ZooModel zooModel = ResNet50.builder().build();
//        ComputationGraph model = (ComputationGraph) zooModel.initPretrained();
        System.out.println(model.summary());

        // 2. 图像预处理（适配ResNet50）
//        Mat image = Imgcodecs.imread("C:\\Users\\liuxia\\Desktop\\宝贝.jpg");
//        Mat processed = new Mat();
//        Imgproc.resize(image, processed, new Size(224, 224));
//        Imgproc.cvtColor(processed, processed, Imgproc.COLOR_BGR2RGB);

//        NativeImageLoader loader = new NativeImageLoader(224, 224, 3, ImageChannelOrder.RGB);
//        INDArray input = loader.asMatrix("C:\\Users\\liuxia\\Desktop\\宝贝.jpg");
//        input.divi(255.0); // 归一化到 [0, 1]

        // 3. 转换为INDArray并标准化
        NativeImageLoader loader = new NativeImageLoader(224, 224, 3);
        INDArray input = loader.asMatrix("C:\\Users\\liuxia\\Desktop\\1720154070611.jpg")
//        INDArray input = loader.asMatrix(processed)
                .divi(255.0);
        String featureLayerName = "bn4a_branch2c"; // 根据实际模型结构修改
        INDArray input1 = model.feedForward(input, false).get(featureLayerName);
        System.out.println(input1.shape()[1]);

        // 创建统计量张量
        INDArray mean = Nd4j.create(new float[]{0.485f, 0.456f, 0.406f}, new int[]{1, 3, 1, 1});
        INDArray std = Nd4j.create(new float[]{0.229f, 0.224f, 0.225f}, new int[]{1, 3, 1, 1});
        input.subi(mean); // 减去均值
        input.divi(std);  // 除以标准差

        // 5. 特征提取
        INDArray features = model.outputSingle(input1);
//        long[] shape = features.shape();
        System.out.println(features.toFloatVector().length);
        for (int i = 0; i < features.toFloatVector().length; i++) {
            System.out.println(features.toFloatVector()[i]);
        }
//        System.out.println("特征向量维度：" + features.shape()[1]); // 输出 2048
    }
}
